import torch
from net.DLinkNet import DLinkNet34
from torch.utils.data import DataLoader
import Dataset
from torchvision.utils import save_image
import tqdm

if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')

    batch_size = 1
    num_worker = 0
    save_path = './test_image'

    test_dataset = Dataset.SegmentationDataset(where='test', seq=None)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False,
                             num_workers=num_worker)

    model = DLinkNet34().to(device)
    model.load_state_dict(torch.load('./model/class12.pt', map_location=device))
    model.eval()

    for i, (x, y, name) in enumerate(tqdm.tqdm(test_loader)):
        x = x.to(device, dtype=torch.float32)

        y_pred = model(x)
        save_image(y_pred, f"{save_path}/{name[0].split('.')[0]}.png")
